from torch import optim
from lib.optimizer import OPTIMIZER


@OPTIMIZER.register('adam')
def get_adam(cfg, model):
    lr = cfg.LR
    wd = cfg.WEIGHT_DECAY
    return optim.Adam(model.parameters(), lr=lr, weight_decay=wd)


@OPTIMIZER.register('sgd')
def get_sgd(cfg, model):
    lr = cfg.LR
    wd = cfg.WEIGHT_DECAY
    mom = cfg.MOMENTUM
    damp = cfg.DAMPENING
    return optim.SGD(model.parameters(), lr=lr, weight_decay=wd, momentum=mom, dampening=damp)
